package com.diao.flink.state;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

/**
 * @author: chenzhidiao
 * @date: 2021/1/19 11:30
 * @description: 测试 KeyedState 的 ValueState
 *              将给定的元数按key进行分组，每组有三个元素是，输出value的平均值
 *
 *              ValueState<T> ：这个状态为每一个 key 保存一个值
 *              value() 获取状态值
 *              update() 更新状态值
 *              clear() 清除状态
 *
 * @version: 1.0
 */
public class KeyedValueStateDemo {
    public static void main(String[] args) throws Exception {
        //初始化编程入口
        StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());

        //获取DataStream
        DataStreamSource<Tuple2<Long, Long>> tuple2DataStreamSource = env.fromElements(Tuple2.of(1L, 3L), Tuple2.of(1L, 5L), Tuple2.of(1L, 7L), Tuple2.of(2L, 4L), Tuple2.of(2L, 2L), Tuple2.of(2L, 5L));

        tuple2DataStreamSource.keyBy(0)
                .flatMap(new CountWindowAverageWithValueState())
                .print();
        env.execute("TestStatefulApi");
    }
}
class CountWindowAverageWithValueState extends RichFlatMapFunction<Tuple2<Long,Long>,Tuple2<Long,Double>>{

    // 用以保存每个 key 出现的次数，以及这个 key 对应的 value 的总值
    // managed keyed state
    // 1. ValueState 保存的是对应的一个 key 的一个状态值
    private ValueState<Tuple2<Long, Long>> countAndSum;


    @Override
    public void open(Configuration parameters) throws Exception{
        //注册一个状态
        ValueStateDescriptor<Tuple2<Long,Long>> descriptor = new ValueStateDescriptor<Tuple2<Long,Long>>(
                "average", //状态的名称
                Types.TUPLE(Types.LONG,Types.LONG));//状态的数据类型

        countAndSum = getRuntimeContext().getState(descriptor);
    }


    @Override
    public void flatMap(Tuple2<Long, Long> value, Collector<Tuple2<Long, Double>> out) throws Exception {

        //获取当前key的状态
        Tuple2<Long, Long> currentState = countAndSum.value();

        //如果状态还未初始化，则进行初始化
        if(currentState == null){
            currentState = Tuple2.of(0L,0L);
        }

        //更新状态的值
        currentState.f0 += 1;
        currentState.f1 += value.f1;

        //更新状态
        countAndSum.update(currentState);

        if(currentState.f0>=3){
            double avg = (double)currentState.f1/currentState.f0;

            //输出
            out.collect(Tuple2.of(value.f0,avg));

            //清空状态值
            countAndSum.clear();
        }

    }
}
